#pragma once

#include <assert.h>
#include <condition_variable>
#include <coroutine>
#include <functional>
#include <optional>
#include "PkgedTask.h"

namespace AsyncLib
{

template<typename T>
using CCoTask = CFuture<T>;

inline std::shared_ptr<CAssocState<void>> g_pCompletedState;

inline void AsyncLibInit(ATID atidMax = ATID_MAIN)
{
	g_pCompletedState = std::make_shared<CAssocState<void>>();
	g_pCompletedState->SetValue();
	GetAsyncLib().Init(atidMax);
}

inline void AsyncLibExit(bool bWaitTasks = true)
{
	GetAsyncLib().Exit(bWaitTasks);
}

// Return a completed Task<void>
[[nodiscard]] inline auto CompletedTask()
{
	return CCoTask<void>(g_pCompletedState);
}

// Return a completed Task<T> with the value
template<typename T>
[[nodiscard]] inline auto TaskFromResult(T&& value)
{
	auto pState = std::make_shared<CAssocState<std::decay_t<T>>>();
	pState->SetValue(std::forward<T>(value));
	return CCoTask<std::decay_t<T>>(std::move(pState));
}

template<typename T, typename Fn, typename... Args>
inline auto PostTaskImpl(T val, Fn&& fn, Args&&... args)
{
	static_assert(std::is_same_v<T, ATID> || std::is_same_v<T, UINT>);
	static_assert(!std::is_null_pointer_v<Fn> || sizeof...(Args) == 0);

	using PkgedTaskT = CPkgedTask<std::decay_t<Fn>, std::decay_t<Args>...>;
	auto pState = std::make_shared<CAssocStateN<typename PkgedTaskT::type, sizeof(PkgedTaskT)>>();
	CAutoRefPtr pPkgedTask(new (pState->GetStorage()) PkgedTaskT(std::forward<Fn>(fn), std::forward<Args>(args)...));
	pPkgedTask->SetState(pState);

	if constexpr (std::is_null_pointer_v<Fn>)
		pState->SuspendFirst(std::is_same_v<T, UINT>);

	if constexpr (std::is_same_v<T, ATID>)
		GetAsyncLib().PostTask(val, std::move(pPkgedTask));
	else
		GetAsyncLib().PostDelayedTask(val, std::move(pPkgedTask));

	return CCoTask<typename PkgedTaskT::type>(std::move(pState));
}

// Post a task in the specified thread
template<typename Fn, typename... Args>
inline auto PostTask(ATID atid, Fn&& fn, Args&&... args)
{
	return PostTaskImpl(atid, std::forward<Fn>(fn), std::forward<Args>(args)...);
}

// Post a task in the thread pool
template<typename Fn, typename... Args>
inline auto PostTask(Fn&& fn, Args&&... args)
{
	return PostTaskImpl(ATID_TP, std::forward<Fn>(fn), std::forward<Args>(args)...);
}

// Post a delayed task in the thread pool
template<typename Fn, typename... Args>
inline auto PostDelayedTask(UINT nMilliseconds, Fn&& fn, Args&&... args)
{
	return PostTaskImpl(nMilliseconds, std::forward<Fn>(fn), std::forward<Args>(args)...);
}

[[nodiscard]] inline bool IsInTaskThread(ATID atid)
{
	return GetAsyncLib().IsInTaskThread(atid);
}

[[nodiscard]] inline bool IsInMainThread()
{
	return IsInTaskThread(ATID_MAIN);
}

[[nodiscard]] inline bool IsInThreadPool()
{
	return GetAsyncLib().IsInThreadPool();
}

// Switch to the specified thread using co_await
[[nodiscard]] inline auto SwitchToTaskThread(ATID atid)
{
	// Do not check if already in the task thread to ensure tasks are in order
	return PostTask(atid, nullptr);
}

[[nodiscard]] inline auto SwitchToMainThread()
{
	return SwitchToTaskThread(ATID_MAIN);
}

// Switch to the thread pool using co_await
[[nodiscard]] inline auto SwitchToThreadPool()
{
	// Do not PostTask if already in the thread pool
	if (IsInThreadPool())
		return TaskFromResult(nullptr);

	return PostTask(nullptr);
}

// Delay the specified time asynchronously
[[nodiscard]] inline auto DelayAsync(UINT nMilliseconds)
{
	return PostDelayedTask(nMilliseconds, nullptr);
}

template<typename Fn, typename... Args>
using InvokeResult = std::invoke_result_t<std::decay_t<Fn>, std::decay_t<Args>...>;

// Overload functions for Fn is a coroutine
template<typename Fn, typename... Args>
	requires IsFuture<InvokeResult<Fn, Args...>>::value
inline InvokeResult<Fn, Args...> PostTask(ATID atid, Fn&& fn, Args&&... args)
{
	auto fnTask = std::bind(std::forward<Fn>(fn), std::forward<Args>(args)...);
	co_await SwitchToTaskThread(atid);
	co_return co_await fnTask();
}

template<typename Fn, typename... Args>
	requires IsFuture<InvokeResult<Fn, Args...>>::value
inline InvokeResult<Fn, Args...> PostTask(Fn&& fn, Args&&... args)
{
	auto fnTask = std::bind(std::forward<Fn>(fn), std::forward<Args>(args)...);
	co_await SwitchToThreadPool();
	co_return co_await fnTask();
}

template<typename Fn, typename... Args>
	requires IsFuture<InvokeResult<Fn, Args...>>::value
inline InvokeResult<Fn, Args...> PostDelayedTask(UINT nMilliseconds, Fn&& fn, Args&&... args)
{
	auto fnTask = std::bind(std::forward<Fn>(fn), std::forward<Args>(args)...);
	co_await DelayAsync(nMilliseconds);
	co_return co_await fnTask();
}

////////////////////////////////////////////////////////////////////////////////

// Wrap a function with a callback as asynchronous
template<typename Ret, typename... FnArgs>
[[nodiscard]] inline auto WrapFnAsAsync(FnArgs&&... fnArgs)
{
	auto pState = std::make_shared<CAssocStateN<Ret, sizeof(CStateManager<Ret>)>>();
	CAutoRefPtr pManager(new (pState->GetStorage()) CStateManager<Ret>());
	pManager->SetState(pState);

	// Complete() must be called in destructor to handle the canceled task
	std::invoke(std::forward<FnArgs>(fnArgs)..., [pManager = std::move(pManager)](auto&&... args)
	{
		pManager->GetState()->SetValue(std::forward<decltype(args)>(args)...);
	});

	return CCoTask<Ret>(std::move(pState));
}

// Wait for the specified object asynchronously, cannot use for Mutex
[[nodiscard]] inline auto WaitForAsync(HANDLE handle, UINT nMilliseconds)
{
	return WrapFnAsAsync<DWORD>(&IAsyncLib::WaitForAsync, GetAsyncLib(), handle, nMilliseconds);
}

// Wait for a task in the main thread to support ATID_MAIN
template<typename T>
[[nodiscard]] inline bool MsgWaitForTask(const CCoTask<T>& task)
{
	assert(IsInMainThread());

#ifdef _WIN32
	auto status = CCoTask<T>::GetState(task)->SetCallback([nThreadId = GetCurrentThreadId()](bool)
	{
		PostThreadMessage(nThreadId, WM_QUIT, 0, 0);
	});

	if (status == CAssocStateBase::Pending)
	{
		MSG msg;
		while (GetMessage(&msg, NULL, 0, 0))
		{
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}
	}
#endif

	return task.Wait();
}

////////////////////////////////////////////////////////////////////////////////

class CWhenAllHandler
{
public:
	CWhenAllHandler(std::shared_ptr<CAssocState<void>> pState) : m_pState(std::move(pState))
	{
	}

	~CWhenAllHandler()
	{
		if (m_pState)
		{
			m_pState->SetValue();
			m_pState->Complete();
		}
	}

	void OnCallback(std::shared_ptr<CAssocStateBase> pState, bool bResult)
	{
		std::unique_lock lock(m_mutex);

		if (m_pState)
		{
			// SetValue() in the destructor when all tasks are completed
			if (bResult && !m_pState->MoveException(*pState))
				return;

			// Stop when any task is canceled or an exception occurs
			m_pState->Complete();
			m_pState = nullptr;
		}
	}

private:
	std::mutex m_mutex;
	std::shared_ptr<CAssocState<void>> m_pState;
};

// Wait for all tasks to complete
template<typename... Ts>
[[nodiscard]] inline auto TaskWhenAll(const CCoTask<Ts>&... tasks)
{
	static_assert(sizeof...(Ts) > 1);
	auto pState = std::make_shared<CAssocState<void>>();
	auto pHandler = std::make_shared<CWhenAllHandler>(pState);

	([&tasks, pHandler]
	{
		auto pState = CCoTask<Ts>::GetState(tasks);
		pState->Forget();

		auto status = pState->SetCallback([pHandler, pState](bool bResult)
		{
			pHandler->OnCallback(pState, bResult);
		});

		if (status != CAssocStateBase::Pending)
			pHandler->OnCallback(pState, status == CAssocStateBase::Completed);
	}(), ...);

	return CCoTask<void>(std::move(pState));
}

////////////////////////////////////////////////////////////////////////////////

class CWhenAnyHandler
{
public:
	CWhenAnyHandler(std::shared_ptr<CAssocState<size_t>> pState) : m_pState(std::move(pState))
	{
	}

	~CWhenAnyHandler()
	{
		assert(!m_pState);
	}

	void OnCallback(std::shared_ptr<CAssocStateBase> pState, bool bResult, size_t nIndex)
	{
		std::unique_lock lock(m_mutex);

		if (m_pState)
		{
			// SetValue() when any task is completed
			if (bResult && !m_pState->MoveException(*pState))
				m_pState->SetValue(nIndex);

			m_pState->Complete();
			m_pState = nullptr;
		}
	}

private:
	std::mutex m_mutex;
	std::shared_ptr<CAssocState<size_t>> m_pState;
};

template<size_t... I, typename... Ts>
inline auto TaskWhenAnyImpl(std::index_sequence<I...>, const CCoTask<Ts>&... tasks)
{
	auto pState = std::make_shared<CAssocState<size_t>>();
	auto pHandler = std::make_shared<CWhenAnyHandler>(pState);

	([&tasks, pHandler]
	{
		auto pState = CCoTask<Ts>::GetState(tasks);
		pState->Forget();

		auto status = pState->SetCallback([pHandler, pState](bool bResult)
		{
			pHandler->OnCallback(pState, bResult, I);
		});

		if (status != CAssocStateBase::Pending)
			pHandler->OnCallback(pState, status == CAssocStateBase::Completed, I);
	}(), ...);

	return CCoTask<size_t>(std::move(pState));
}

// Wait for any task to complete
template<typename... Ts>
[[nodiscard]] inline auto TaskWhenAny(const CCoTask<Ts>&... tasks)
{
	static_assert(sizeof...(Ts) > 1);
	return TaskWhenAnyImpl(std::index_sequence_for<Ts...>{}, tasks...);
}

////////////////////////////////////////////////////////////////////////////////

template<bool bUI>
class CTimerT
{
public:
	CTimerT(std::function<void()>&& fnCallback) : m_pTimerTask(GetAsyncLib().CreateTimerTask(bUI))
	{
		assert(fnCallback);

		if (m_pTimerTask)
			m_pTimerTask->SetCallback(std::move(fnCallback));
	}

	~CTimerT()
	{
		if (m_pTimerTask)
			m_pTimerTask->Close();
	}

	CTimerT(CTimerT&&) = default;
	CTimerT& operator=(CTimerT&&) = default;

	void Start(UINT nElapse) const
	{
		if (m_pTimerTask)
			m_pTimerTask->Start(nElapse);
	}

	void Stop() const
	{
		if (m_pTimerTask)
			m_pTimerTask->Stop();
	}

private:
	std::shared_ptr<ITimerTask> m_pTimerTask;
};

} // namespace AsyncLib

namespace AL
{

template<typename T = void>
using Task = AsyncLib::CCoTask<T>;

using AsyncLib::ATID;
using AsyncLib::ATID_MAIN;
using AsyncLib::AsyncLibInit;
using AsyncLib::AsyncLibExit;

using AsyncLib::CompletedTask;
using AsyncLib::TaskFromResult;

using AsyncLib::PostTask;
using AsyncLib::PostDelayedTask;

using AsyncLib::IsInMainThread;
using AsyncLib::IsInTaskThread;
using AsyncLib::IsInThreadPool;

using AsyncLib::SwitchToMainThread;
using AsyncLib::SwitchToTaskThread;
using AsyncLib::SwitchToThreadPool;
using AsyncLib::DelayAsync;
using AsyncLib::WaitForAsync;
using AsyncLib::MsgWaitForTask;
using AsyncLib::WrapFnAsAsync;

using AsyncLib::TaskWhenAll;
using AsyncLib::TaskWhenAny;

using CUiTimer = AsyncLib::CTimerT<true>;
using CTpTimer = AsyncLib::CTimerT<false>;

}

#include "AsyncMutex.h"
#if __has_include(<windows.h>)
#include "AsyncFile.h"
#endif
